#!/usr/bin/env node
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
    CallToolRequestSchema,
    ErrorCode,
    ListResourcesRequestSchema,
    ListResourceTemplatesRequestSchema,
    ListToolsRequestSchema,
    McpError,
    ReadResourceRequestSchema,
} from '@modelcontextprotocol/sdk/types.js';
import { MemoryStorage } from './storage.js';
import { MemoryAnalyzer } from './analyzer.js';
import { MemoryType } from './types.js';
import path from 'path';
import os from 'os';

/**
 * AI记忆服务器类
 * 实现MCP服务器，提供AI记忆相关的工具和资源
 */
class AIMemoryServer {
    private server: Server;
    private storage: MemoryStorage;

    constructor() {
        // 初始化MCP服务器
        this.server = new Server(
            {
                name: 'ai-memory-server',
                version: '0.1.0',
            },
            {
                capabilities: {
                    resources: {},
                    tools: {},
                },
            }
        );

        // 初始化存储
        const dbPath = path.join(os.homedir(), '.ai-memory.db');
        this.storage = new MemoryStorage(dbPath);

        // 设置请求处理器
        this.setupResourceHandlers();
        this.setupToolHandlers();

        // 错误处理
        this.server.onerror = (error) => console.error('[MCP Error]', error);
        process.on('SIGINT', async () => {
            await this.storage.close();
            await this.server.close();
            process.exit(0);
        });
    }

    /**
     * 设置资源处理器
     */
    private setupResourceHandlers() {
        // 列出可用资源
        this.server.setRequestHandler(ListResourcesRequestSchema, async () => ({
            resources: [
                {
                    uri: 'ai-memory://error-patterns',
                    name: '错误模式',
                    mimeType: 'application/json',
                    description: '常见错误模式及其解决方案',
                },
                {
                    uri: 'ai-memory://experience-summary',
                    name: '经验总结',
                    mimeType: 'text/markdown',
                    description: 'AI经验总结',
                },
                {
                    uri: 'ai-memory://thought-analysis',
                    name: '思考分析',
                    mimeType: 'text/markdown',
                    description: 'AI思考过程分析',
                },
            ],
        }));

        // 列出资源模板
        this.server.setRequestHandler(
            ListResourceTemplatesRequestSchema,
            async () => ({
                resourceTemplates: [
                    {
                        uriTemplate: 'ai-memory://experience-summary/{topic}',
                        name: '特定主题的经验总结',
                        mimeType: 'text/markdown',
                        description: '获取特定主题的AI经验总结',
                    },
                ],
            })
        );

        // 读取资源
        this.server.setRequestHandler(
            ReadResourceRequestSchema,
            async (request) => {
                const uri = request.params.uri;

                // 错误模式资源
                if (uri === 'ai-memory://error-patterns') {
                    const patterns = await this.storage.getErrorPatterns();
                    const report = MemoryAnalyzer.analyzeErrorPatterns(patterns);

                    return {
                        contents: [
                            {
                                uri,
                                mimeType: 'text/markdown',
                                text: report,
                            },
                        ],
                    };
                }

                // 经验总结资源
                else if (uri === 'ai-memory://experience-summary') {
                    const experiences = await this.storage.retrieveMemories(
                        MemoryType.EXPERIENCE,
                        100
                    );
                    const summary = MemoryAnalyzer.generateExperienceSummary(experiences);

                    return {
                        contents: [
                            {
                                uri,
                                mimeType: 'text/markdown',
                                text: summary,
                            },
                        ],
                    };
                }

                // 特定主题的经验总结
                else if (uri.startsWith('ai-memory://experience-summary/')) {
                    const topic = uri.replace('ai-memory://experience-summary/', '');
                    const experiences = await this.storage.retrieveMemories(
                        MemoryType.EXPERIENCE,
                        100,
                        0,
                        topic
                    );
                    const summary = MemoryAnalyzer.generateExperienceSummary(experiences, topic);

                    return {
                        contents: [
                            {
                                uri,
                                mimeType: 'text/markdown',
                                text: summary,
                            },
                        ],
                    };
                }

                // 思考分析资源
                else if (uri === 'ai-memory://thought-analysis') {
                    const thoughts = await this.storage.retrieveMemories(
                        MemoryType.THOUGHT,
                        100
                    );
                    const analysis = MemoryAnalyzer.analyzeThoughts(thoughts);

                    return {
                        contents: [
                            {
                                uri,
                                mimeType: 'text/markdown',
                                text: analysis,
                            },
                        ],
                    };
                }

                // 未知资源
                else {
                    throw new McpError(
                        ErrorCode.InvalidRequest,
                        `未知资源: ${uri}`
                    );
                }
            }
        );
    }

    /**
     * 设置工具处理器
     */
    private setupToolHandlers() {
        // 列出可用工具
        this.server.setRequestHandler(ListToolsRequestSchema, async () => ({
            tools: [
                {
                    name: 'record_error',
                    description: '记录错误',
                    inputSchema: {
                        type: 'object',
                        properties: {
                            content: {
                                type: 'string',
                                description: '错误内容',
                            },
                            context: {
                                type: 'string',
                                description: '错误发生的上下文',
                            },
                            tags: {
                                type: 'string',
                                description: '标签，用逗号分隔',
                            },
                        },
                        required: ['content'],
                    },
                },
                {
                    name: 'record_experience',
                    description: '记录经验',
                    inputSchema: {
                        type: 'object',
                        properties: {
                            content: {
                                type: 'string',
                                description: '经验内容',
                            },
                            context: {
                                type: 'string',
                                description: '经验的上下文',
                            },
                            tags: {
                                type: 'string',
                                description: '标签，用逗号分隔',
                            },
                        },
                        required: ['content'],
                    },
                },
                {
                    name: 'record_thought',
                    description: '记录思考',
                    inputSchema: {
                        type: 'object',
                        properties: {
                            content: {
                                type: 'string',
                                description: '思考内容',
                            },
                            context: {
                                type: 'string',
                                description: '思考的上下文',
                            },
                            tags: {
                                type: 'string',
                                description: '标签，用逗号分隔',
                            },
                        },
                        required: ['content'],
                    },
                },
                {
                    name: 'search_memories',
                    description: '搜索记忆',
                    inputSchema: {
                        type: 'object',
                        properties: {
                            type: {
                                type: 'string',
                                enum: ['ERROR', 'EXPERIENCE', 'THOUGHT', 'all'],
                                description: '记忆类型',
                            },
                            searchText: {
                                type: 'string',
                                description: '搜索文本',
                            },
                            tags: {
                                type: 'string',
                                description: '标签过滤，用逗号分隔',
                            },
                            limit: {
                                type: 'number',
                                description: '返回结果数量限制',
                            },
                        },
                    },
                },
                {
                    name: 'update_error_solution',
                    description: '更新错误解决方案',
                    inputSchema: {
                        type: 'object',
                        properties: {
                            patternId: {
                                type: 'number',
                                description: '错误模式ID',
                            },
                            solution: {
                                type: 'string',
                                description: '解决方案',
                            },
                        },
                        required: ['patternId', 'solution'],
                    },
                },
            ],
        }));

        // 调用工具
        this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
            const { name, arguments: toolArgs } = request.params;

            try {
                // 记录错误
                if (name === 'record_error') {
                    if (!toolArgs || !toolArgs.content) {
                        throw new McpError(
                            ErrorCode.InvalidParams,
                            '缺少必要参数: content'
                        );
                    }

                    const id = await this.storage.storeMemory(
                        MemoryType.ERROR,
                        toolArgs.content as string,
                        toolArgs.context as string | undefined,
                        toolArgs.tags as string | undefined
                    );

                    return {
                        content: [
                            {
                                type: 'text',
                                text: `错误已记录，ID: ${id}`,
                            },
                        ],
                    };
                }

                // 记录经验
                else if (name === 'record_experience') {
                    if (!toolArgs || !toolArgs.content) {
                        throw new McpError(
                            ErrorCode.InvalidParams,
                            '缺少必要参数: content'
                        );
                    }

                    const id = await this.storage.storeMemory(
                        MemoryType.EXPERIENCE,
                        toolArgs.content as string,
                        toolArgs.context as string | undefined,
                        toolArgs.tags as string | undefined
                    );

                    return {
                        content: [
                            {
                                type: 'text',
                                text: `经验已记录，ID: ${id}`,
                            },
                        ],
                    };
                }

                // 记录思考
                else if (name === 'record_thought') {
                    if (!toolArgs || !toolArgs.content) {
                        throw new McpError(
                            ErrorCode.InvalidParams,
                            '缺少必要参数: content'
                        );
                    }

                    const id = await this.storage.storeMemory(
                        MemoryType.THOUGHT,
                        toolArgs.content as string,
                        toolArgs.context as string | undefined,
                        toolArgs.tags as string | undefined
                    );

                    return {
                        content: [
                            {
                                type: 'text',
                                text: `思考已记录，ID: ${id}`,
                            },
                        ],
                    };
                }

                // 搜索记忆
                else if (name === 'search_memories') {
                    const searchArgs = toolArgs || {};
                    const type = searchArgs.type === 'all' ? undefined :
                        searchArgs.type ? MemoryType[searchArgs.type as keyof typeof MemoryType] : undefined;
                    const limit = searchArgs.limit ? Number(searchArgs.limit) : 10;

                    const memories = await this.storage.retrieveMemories(
                        type,
                        limit,
                        0,
                        searchArgs.tags as string | undefined,
                        searchArgs.searchText as string | undefined
                    );

                    return {
                        content: [
                            {
                                type: 'text',
                                text: `找到 ${memories.length} 条记忆:\n\n${JSON.stringify(memories, null, 2)}`,
                            },
                        ],
                    };
                }

                // 更新错误解决方案
                else if (name === 'update_error_solution') {
                    if (!toolArgs || !toolArgs.patternId || !toolArgs.solution) {
                        throw new McpError(
                            ErrorCode.InvalidParams,
                            '缺少必要参数: patternId 或 solution'
                        );
                    }

                    const success = await this.storage.updateErrorSolution(
                        toolArgs.patternId as number,
                        toolArgs.solution as string
                    );

                    return {
                        content: [
                            {
                                type: 'text',
                                text: success
                                    ? '解决方案已更新'
                                    : '更新失败，可能是因为找不到指定的错误模式',
                            },
                        ],
                    };
                }

                // 未知工具
                else {
                    throw new McpError(
                        ErrorCode.MethodNotFound,
                        `未知工具: ${name}`
                    );
                }
            } catch (error) {
                if (error instanceof McpError) {
                    throw error;
                }

                console.error(`工具执行错误:`, error);
                throw new McpError(
                    ErrorCode.InternalError,
                    `工具执行错误: ${(error as Error).message}`
                );
            }
        });
    }

    /**
     * 运行服务器
     */
    async run() {
        // 初始化存储
        await this.storage.initialize();

        // 连接到传输层
        const transport = new StdioServerTransport();
        await this.server.connect(transport);
        console.error('AI Memory MCP 服务器已启动');
    }
}

// 创建并运行服务器
const server = new AIMemoryServer();
server.run().catch(console.error);